#include <iostream>
#include <opencv2/core/core.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/highgui/highgui.hpp>
#include "vortex/engine/simple_infer_engine.h"
#include "vortex/layers/softmax_prob.h"


namespace vortex
{
    class Vanillanet : public SimpleInferEngine
    {
    public:
        Vanillanet(const std::string& engine_path, uint32_t num_classes = 1000)
        {
            BlobInfo input_info = {
                "input", 1, 3, 224, 224
            };

            BlobInfo output_info = {
                "output", 1, num_classes, 1, 1
            };

            this->LoadEngine(engine_path, input_info, output_info);
        }

        void Preprocess(cv::Mat& image)
        {
            float mean[3] = {0.485, 0.456, 0.406};
            float std[3] = {0.229, 0.224, 0.225};
            // resize, bgr-rgb, normalization with imagenet mean/std
            cv::Mat temp;
            uint32_t width = m_InputInfo.width;
            uint32_t height = m_InputInfo.height;
            // uint32_t channels = m_InputInfo.channels;
            cv::resize(image, temp, cv::Size(width, height));
            uint32_t image_area = width * height;
            // uint32_t numel = image_area * channels;

            float* input_buffer = m_InputBlob->dataCpu;
            float* pBlue = input_buffer;
            float* pGreen = input_buffer + image_area;
            float* pRed = input_buffer + image_area * 2;
            unsigned char* pImage = temp.data;
            for (uint32_t i = 0; i < image_area; ++i)
            {
                pRed[i] = (pImage[3 * i + 0] / 255.0 - mean[0]) / std[0];
                pGreen[i] = (pImage[3 * i + 1] / 255.0 - mean[1]) / std[1];
                pBlue[i] = (pImage[3 * i + 2] / 255.0 - mean[2]) / std[2];
            }
        }
    };
}


int main()
{
    std::string engine_path("../vanillanet_5.engine");
    vortex::Vanillanet model(engine_path);
    std::string image_path("../../../data/person.jpg");
    cv::Mat image = cv::imread(image_path);
    vortex::SoftmaxProb softmax(1000);

    std::vector<float> outputs;
    model.Infer(image, outputs);
    std::vector<float> probs;
    softmax.Forward(outputs, probs);

    for (int i = 0; i < 10; ++i)
    {
        std::cout << probs[i] << std::endl;
    }

    std::cout << "Finished\n";
    return 0;
}
